import tkinter as tk
from tkinter import ttk, messagebox, filedialog
import torch
import torch.nn.functional as F
import numpy as np
from torchvision import transforms
from model import MOE_Model
from PIL import Image, ImageDraw
import os
from config import ensure_directories, RUNS_DIR  # 导入配置


class DigitRecognizer:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("手写数字识别")

        # 设置主题色
        self.colors = {
            "primary": "#2196F3",  # 主色调
            "secondary": "#FFC107",  # 强调色
            "success": "#4CAF50",  # 成功色
            "danger": "#f44336",  # 警告色
            "bg": "#F5F5F5",  # 背景色
            "text": "#212121",  # 文字色
            "light": "#FFFFFF",  # 亮色
            "dark": "#333333",  # 暗色
        }

        # 设置样式
        style = ttk.Style()

        # 主按钮样式 (蓝色)
        style.configure(
            "Primary.TButton",
            background=self.colors["primary"],
            foreground=self.colors["light"],
            font=("Microsoft YaHei UI", 10),
        )

        # 强调按钮样式 (橙色)
        style.configure(
            "Accent.TButton",
            background=self.colors["secondary"],
            foreground=self.colors["dark"],
            font=("Microsoft YaHei UI", 10, "bold"),
        )

        # 危险按钮样式 (红色)
        style.configure(
            "Danger.TButton",
            background=self.colors["danger"],
            foreground=self.colors["light"],
            font=("Microsoft YaHei UI", 10),
        )

        # 普通按钮样式
        style.configure("TButton", font=("Microsoft YaHei UI", 10))

        # 初始化参数
        self.canvas_size = 400
        self.pen_size = 25
        self.drawing = False
        self.last_x = None
        self.last_y = None
        self.draw_history = []  # 用于存储绘图操作
        self.current_stroke = []  # 用于存储当前笔画
        self.zoom_level = 1.0  # 画布缩放级别
        self.show_grid = tk.BooleanVar(value=False)  # 添加网格显示控制

        # 添加画笔设置
        self.pen_color = "black"
        self.pen_styles = {"round": tk.ROUND, "square": tk.PROJECTING, "butt": tk.BUTT}
        self.current_pen_style = tk.ROUND

        # 初始化窗口和组件
        self.setup_window()
        self.init_ui()

        # 初始化设备和模型
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = None
        self.disable_recognition()  # 初始禁用所有操作按钮

        # 更新状态栏提示
        self.status_label.config(text="请选择并加载模型")

        ensure_directories()  # 确保目录存在

    def setup_window(self):
        """设置窗口属性"""
        # 设置窗口大小和位置
        window_width = 800
        window_height = 600
        screen_width = self.root.winfo_screenwidth()
        screen_height = self.root.winfo_screenheight()
        center_x = int(screen_width / 2 - window_width / 2)
        center_y = int(screen_height / 2 - window_height / 2)

        self.root.geometry(f"{window_width}x{window_height}+{center_x}+{center_y}")
        self.root.resizable(False, False)
        self.root.configure(bg=self.colors["bg"])

    def init_ui(self):
        """初始化用户界面"""
        # 创建主框架
        main_frame = ttk.Frame(self.root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=20, pady=20)

        # 左侧面板 - 绘图区域 (占60%宽度)
        left_panel = ttk.Frame(main_frame)
        left_panel.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=(0, 10))

        # 右侧面板 - 控制和显示区域 (占40%宽度)
        right_panel = ttk.Frame(main_frame)
        right_panel.pack(side=tk.LEFT, fill=tk.BOTH, padx=(10, 0))

        # === 左侧面板布局 ===
        # 标题
        title = tk.Label(
            left_panel,
            text="手写数字识别",
            font=("Microsoft YaHei UI", 24, "bold"),
            bg=self.colors["bg"],
            fg=self.colors["text"],
        )
        title.pack(pady=(0, 20))

        # 画布区域
        canvas_frame = ttk.Frame(left_panel)
        canvas_frame.pack(pady=(0, 20))

        self.setup_canvas(canvas_frame)

        # === 右侧面板布局 ===
        # 模型设置区域
        model_frame = ttk.LabelFrame(right_panel, text="模型设置", padding=10)
        model_frame.pack(fill=tk.X, pady=(0, 10))

        model_path_frame = ttk.Frame(model_frame)
        model_path_frame.pack(fill=tk.X, pady=(0, 10))

        self.model_path_var = tk.StringVar()
        model_path_entry = ttk.Entry(
            model_path_frame, textvariable=self.model_path_var, state="readonly"
        )
        model_path_entry.pack(fill=tk.X, padx=5)

        btn_frame = ttk.Frame(model_frame)
        btn_frame.pack(fill=tk.X)

        ttk.Button(
            btn_frame,
            text="选择模型",
            command=self.choose_model,
            width=15,
            style="Primary.TButton",
        ).pack(side=tk.LEFT, padx=5, expand=True)

        ttk.Button(
            btn_frame,
            text="加载模型",
            command=self.load_selected_model,
            width=15,
            style="Primary.TButton",
        ).pack(side=tk.LEFT, padx=5, expand=True)

        # 画笔设置区域
        pen_frame = ttk.LabelFrame(right_panel, text="画笔设置", padding=10)
        pen_frame.pack(fill=tk.X, pady=10)

        # 画笔大小
        size_frame = ttk.Frame(pen_frame)
        size_frame.pack(fill=tk.X, pady=2)
        ttk.Label(size_frame, text="大小:").pack(side=tk.LEFT, padx=5)
        self.pen_scale = ttk.Scale(
            size_frame,
            from_=5,
            to=50,
            orient=tk.HORIZONTAL,
            value=self.pen_size,
            command=self.update_pen_size,
        )
        self.pen_scale.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5)
        self.size_label = ttk.Label(size_frame, text=str(self.pen_size))
        self.size_label.pack(side=tk.LEFT, padx=5)

        # 画笔颜色
        color_frame = ttk.Frame(pen_frame)
        color_frame.pack(fill=tk.X, pady=2)
        ttk.Label(color_frame, text="颜色:").pack(side=tk.LEFT, padx=5)

        self.color_var = tk.StringVar(value="black")
        colors = [
            ("黑色", "black"),
            ("蓝色", "blue"),
            ("红色", "red"),
            ("绿色", "green"),
        ]
        for name, color in colors:
            ttk.Radiobutton(
                color_frame,
                text=name,
                variable=self.color_var,
                value=color,
                command=lambda c=color: self.change_pen_color(c),
            ).pack(side=tk.LEFT, padx=5)

        # 画笔样式
        style_frame = ttk.Frame(pen_frame)
        style_frame.pack(fill=tk.X, pady=2)
        ttk.Label(style_frame, text="笔尖:").pack(side=tk.LEFT, padx=5)

        self.style_var = tk.StringVar(value="round")
        styles = [("圆形", "round"), ("方形", "square"), ("平直", "butt")]
        for name, style in styles:
            ttk.Radiobutton(
                style_frame,
                text=name,
                variable=self.style_var,
                value=style,
                command=lambda s=style: self.change_pen_style(s),
            ).pack(side=tk.LEFT, padx=5)

        # 控制面板
        control_frame = ttk.LabelFrame(right_panel, text="操作面板", padding=10)
        control_frame.pack(fill=tk.X, pady=10)

        # 主要操作按钮
        main_btn_frame = ttk.Frame(control_frame)
        main_btn_frame.pack(fill=tk.X, pady=(0, 10))

        ttk.Button(
            main_btn_frame, text="识别", command=self.predict, style="Accent.TButton"
        ).pack(side=tk.LEFT, expand=True, padx=5)

        ttk.Button(
            main_btn_frame, text="清除", command=self.clear, style="Danger.TButton"
        ).pack(side=tk.LEFT, expand=True, padx=5)

        # 辅助操作按钮
        sub_btn_frame = ttk.Frame(control_frame)
        sub_btn_frame.pack(fill=tk.X)

        ttk.Button(sub_btn_frame, text="撤销", command=self.undo).pack(
            side=tk.LEFT, expand=True, padx=5
        )

        ttk.Button(
            sub_btn_frame, text="清除最后一笔", command=self.clear_last_stroke
        ).pack(side=tk.LEFT, expand=True, padx=5)

        # 结果显示区域
        result_frame = ttk.LabelFrame(right_panel, text="识别结果", padding=10)
        result_frame.pack(fill=tk.X, pady=10)

        self.result_label = tk.Label(
            result_frame,
            text="等待识别...",
            font=("Microsoft YaHei UI", 14),
            bg=self.colors["light"],
            fg=self.colors["text"],
            height=2,
        )
        self.result_label.pack(fill=tk.X)

        # 状态栏
        self.status_label = tk.Label(
            self.root,
            text="就绪",
            font=("Microsoft YaHei UI", 10),
            bg=self.colors["primary"],
            fg=self.colors["light"],
            anchor="w",
            padx=10,
            pady=5,
        )
        self.status_label.pack(side=tk.BOTTOM, fill=tk.X)

        # 绑定快捷键
        self.root.bind("<Control-z>", lambda e: self.undo())
        self.root.bind("<Return>", lambda e: self.predict())
        self.root.bind("<Control-plus>", lambda e: self.zoom_canvas(1.2))  # Ctrl++ 放大
        self.root.bind(
            "<Control-minus>", lambda e: self.zoom_canvas(0.8)
        )  # Ctrl+- 缩小
        self.root.bind(
            "<Control-0>", lambda e: self.zoom_canvas(reset=True)
        )  # Ctrl+0 重置缩放
        self.root.bind(
            "<Control-b>", lambda e: self.color_var.set("black")
        )  # Ctrl+B 黑色画笔
        self.root.bind(
            "<Control-r>", lambda e: self.color_var.set("red")
        )  # Ctrl+R 红色画笔
        self.root.bind(
            "<Control-l>", lambda e: self.clear_last_stroke()
        )  # Ctrl+L 清除最后一笔

        # 绑定鼠标事件
        self.canvas.bind("<Button-1>", self.start_drawing)
        self.canvas.bind("<B1-Motion>", self.draw_line)
        self.canvas.bind("<ButtonRelease-1>", self.stop_drawing)

    def update_pen_size(self, value):
        """更新画笔大小"""
        self.pen_size = int(float(value))
        self.size_label.config(text=str(self.pen_size))
        self.status_label.config(text=f"画笔大小已更改为: {self.pen_size}")

    def undo(self):
        """撤销上一步绘制"""
        if self.draw_history:
            self.draw_history.pop()
            self.redraw_canvas()
            self.status_label.config(text="已撤销上一步操作")
        else:
            self.status_label.config(text="没有可撤销的操作")

    def load_model(self):
        """加载模型"""
        try:
            self.status_label.config(text="正在加载模型...")
            self.root.update()

            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

            # 优先加载最新的JIT模型
            latest_jit_path = "runs/best_model_jit.pth"

            if os.path.exists(latest_jit_path):
                try:
                    self.model = torch.jit.load(
                        latest_jit_path, map_location=self.device
                    )
                    print("已加载最新JIT模型")
                except Exception as e:
                    print(f"加载JIT模型失败: {e}")
                    # 回退到加载常规模型
                    self.load_regular_model()
            else:
                print("未找到JIT模型，加载常规模型...")
                self.load_regular_model()

            self.model.eval()

            # 预热模型
            print("预热模型...")
            dummy_input = torch.randn(1, 1, 28, 28).to(self.device)
            with torch.no_grad():
                _ = self.model(dummy_input)

            device_name = "GPU" if self.device.type == "cuda" else "CPU"
            self.status_label.config(text=f"模型加载完成，使用{device_name}")

        except Exception as e:
            error_msg = f"模型加载失败: {str(e)}"
            print(error_msg)
            self.status_label.config(text=error_msg)

            # 显示错误对话框，提供重试选项
            retry = messagebox.askretrycancel(
                "错误", f"模型加载失败:\n{str(e)}\n\n是否重试加载模型？", icon="error"
            )

            if retry:
                # 如果用户选择重试，递归调用加载函数
                self.load_model()
            else:
                # 如果用户取消，禁用识别按钮
                self.disable_recognition()
                self.status_label.config(text="模型未加载，识别功能已禁用")

    def load_regular_model(self):
        """加载常规模型的回退方法"""
        self.model = MOE_Model(num_experts=5).to(self.device)
        checkpoint = torch.load(
            "runs/best_model.pth", map_location=self.device, weights_only=True
        )
        state_dict = checkpoint.get("model_state_dict", checkpoint)
        self.model.load_state_dict(state_dict)

    def start_drawing(self, event):
        """开始绘制"""
        if not self.model:  # 如果模型未加载，不允许绘制
            self.status_label.config(text="请先加载模型...")
            return
        
        self.drawing = True
        self.last_x = event.x
        self.last_y = event.y
        self.current_stroke = [(event.x, event.y)]
        self.status_label.config(text="正在绘制...")

    def draw_line(self, event):
        """优化的绘制线条函数"""
        if self.drawing and self.last_x and self.last_y:
            # 计算当前点与上一点的距离
            distance = (
                (event.x - self.last_x) ** 2 + (event.y - self.last_y) ** 2
            ) ** 0.5

            # 如果距离太大，创建平滑的中间点
            if distance > 20:
                steps = int(distance / 10)
                for i in range(1, steps):
                    t = i / steps
                    x = self.last_x + (event.x - self.last_x) * t
                    y = self.last_y + (event.y - self.last_y) * t
                    # 在画布上绘制
                    self.canvas.create_line(
                        self.last_x,
                        self.last_y,
                        x,
                        y,
                        width=self.pen_size,
                        fill=self.pen_color,
                        capstyle=self.current_pen_style,
                        smooth=True,
                    )
                    # 在PIL图像上绘制
                    self.draw.line(
                        [self.last_x, self.last_y, x, y],
                        fill="black",
                        width=self.pen_size,
                    )
                    self.last_x, self.last_y = x, y
                    self.current_stroke.append((x, y))

            # 绘制到当前点
            self.canvas.create_line(
                self.last_x,
                self.last_y,
                event.x,
                event.y,
                width=self.pen_size,
                fill=self.pen_color,
                capstyle=self.current_pen_style,
                smooth=True,
            )
            self.draw.line(
                [self.last_x, self.last_y, event.x, event.y],
                fill="black",
                width=self.pen_size,
            )

            self.current_stroke.append((event.x, event.y))
            self.last_x = event.x
            self.last_y = event.y

    def stop_drawing(self, event):
        """停止绘制"""
        if self.drawing:
            self.draw_history.append(self.current_stroke)
            self.current_stroke = []
            self.status_label.config(text="就绪")
        self.drawing = False
        self.last_x = None
        self.last_y = None

    def clear(self):
        """清除画布"""
        self.canvas.delete("all")
        self.image = Image.new("L", (self.canvas_size, self.canvas_size), "white")
        self.draw = ImageDraw.Draw(self.image)
        self.draw_history = []
        self.result_label.config(text="等待识别...")
        self.status_label.config(text="画布已清空")

    def predict(self):
        """预测函数"""
        try:
            self.status_label.config(text="正在识别...")
            self.root.update()

            # 预处理图像
            img = self.preprocess_image()

            # 预测
            with torch.no_grad():
                output = self.model(img)
                probs = F.softmax(output, dim=1)
                pred = torch.argmax(output, dim=1).item()
                conf = probs[0][pred].item() * 100

            # 显示结果
            self.result_label.config(
                text=f"预测结果: {pred}\n置信度: {conf:.2f}%", fg=self.colors["success"]
            )
            self.status_label.config(text="识别完成")

        except Exception as e:
            self.status_label.config(text=f"识别失败: {str(e)}")
            self.result_label.config(text="识别失败", fg=self.colors["danger"])

    def preprocess_image(self):
        """图像预处理"""
        img = self.image.resize((28, 28), Image.LANCZOS)
        img = np.array(img)
        img = 255 - img  # 反转颜色

        # 转换为tensor
        img_tensor = torch.FloatTensor(img).unsqueeze(0).unsqueeze(0)
        img_tensor = img_tensor / 255.0
        img_tensor = transforms.Normalize((0.1307,), (0.3081,))(img_tensor)
        return img_tensor.to(self.device)

    def toggle_grid(self):
        """切换网格显示"""
        if self.show_grid.get():
            # 绘制网格
            for i in range(0, self.canvas_size, 50):
                self.canvas.create_line(
                    i, 0, i, self.canvas_size, fill="#CCCCCC", dash=(1, 2)
                )
                self.canvas.create_line(
                    0, i, self.canvas_size, i, fill="#CCCCCC", dash=(1, 2)
                )
        else:
            # 清除网格（保留绘图）
            self.redraw_canvas()

    def redraw_canvas(self):
        """重绘画布内容"""
        self.canvas.delete("all")

        # 重绘所有笔画
        for stroke in self.draw_history:
            for i in range(len(stroke) - 1):
                x1, y1 = stroke[i]
                x2, y2 = stroke[i + 1]
                self.canvas.create_line(
                    x1,
                    y1,
                    x2,
                    y2,
                    width=self.pen_size,
                    fill=self.pen_color,
                    capstyle=self.current_pen_style,
                    smooth=True,
                )

        # 如果启用了网格，则绘制网格
        if self.show_grid.get():
            for i in range(0, self.canvas_size, 50):
                self.canvas.create_line(
                    i, 0, i, self.canvas_size, fill="#CCCCCC", dash=(1, 2)
                )
                self.canvas.create_line(
                    0, i, self.canvas_size, i, fill="#CCCCCC", dash=(1, 2)
                )

    def zoom_canvas(self, factor=1.0, reset=False):
        """缩放画布"""
        if reset:
            self.zoom_level = 1.0
        else:
            self.zoom_level *= factor
            self.zoom_level = max(0.5, min(2.0, self.zoom_level))  # 限制缩放范围

        # 调整画笔大小
        base_pen_size = 25  # 基础画笔大小
        self.pen_size = int(base_pen_size * self.zoom_level)
        self.pen_scale.set(self.pen_size)

        # 重绘画布
        self.redraw_canvas()

    def run(self):
        """运行程序"""
        self.root.mainloop()

    def change_pen_color(self, color):
        """更改画笔颜色"""
        self.pen_color = color
        self.status_label.config(text=f"画笔颜色已更改为: {color}")

    def change_pen_style(self, style):
        """更改画笔样式"""
        self.current_pen_style = self.pen_styles[style]
        style_names = {"round": "圆形", "square": "方形", "butt": "平直"}
        self.status_label.config(text=f"笔尖样式已更改为: {style_names[style]}")

    def clear_last_stroke(self):
        """清除最后一笔"""
        if self.draw_history:
            self.draw_history.pop()
            self.redraw_canvas()
            self.image = Image.new("L", (self.canvas_size, self.canvas_size), "white")
            self.draw = ImageDraw.Draw(self.image)
            for stroke in self.draw_history:
                for i in range(len(stroke) - 1):
                    x1, y1 = stroke[i]
                    x2, y2 = stroke[i + 1]
                    self.draw.line([x1, y1, x2, y2], fill="black", width=self.pen_size)
            self.status_label.config(text="已清除最后一笔")
        else:
            self.status_label.config(text="没有可清除的笔画")

    def disable_recognition(self):
        """禁用识别功能"""
        self.update_buttons_state("disabled")
        self.result_label.config(text="请先加载模型...")
        self.status_label.config(text="模型未加载，功能已禁用")

    def enable_recognition(self):
        """启用识别功能"""
        self.update_buttons_state("normal")
        self.result_label.config(text="等待识别...")
        self.status_label.config(text="模型已加载，可以开始使用")

    def load_selected_model(self):
        """加载选中的模型"""
        model_path = self.model_path_var.get()
        if not model_path:
            messagebox.showwarning("提示", "请先选择模型文件")
            return

        try:
            self.status_label.config(text="正在加载模型...")
            self.root.update()

            # 根据文件名判断是否是JIT模型
            is_jit = "_jit.pth" in model_path.lower()

            # 加载模型
            if is_jit:
                self.model = torch.jit.load(model_path, map_location=self.device)
            else:
                self.model = MOE_Model(num_experts=5).to(self.device)
                checkpoint = torch.load(
                    model_path, map_location=self.device, weights_only=False
                )
                self.model.load_state_dict(checkpoint["model_state_dict"])

            self.model.eval()

            # 预热模型
            dummy_input = torch.randn(1, 1, 28, 28).to(self.device)
            with torch.no_grad():
                _ = self.model(dummy_input)

            device_name = "GPU" if self.device.type == "cuda" else "CPU"
            self.status_label.config(text=f"模型已加载，使用{device_name}")
            self.enable_recognition()

        except Exception as e:
            self.status_label.config(text=f"模型加载失败: {str(e)}")
            self.disable_recognition()

    def choose_model(self):
        """选择模型文件"""
        filetypes = [
            ("所有模型文件", "*.pth"),
            ("JIT模型", "*_jit.pth"),
            ("常规模型", "best_model.pth"),
            ("所有文件", "*.*"),
        ]

        model_path = filedialog.askopenfilename(
            title="选择模型文件", initialdir=RUNS_DIR, filetypes=filetypes
        )

        if model_path:
            self.model_path_var.set(model_path)
            if messagebox.askyesno("确认", "是否立即加载选中的模型？"):
                self.load_selected_model()

    def setup_canvas(self, canvas_frame):
        """设置画布和绘图相关功能"""
        # 创建画布边框
        canvas_border = tk.Frame(
            canvas_frame, bg=self.colors["primary"], padx=2, pady=2
        )
        canvas_border.pack()

        # 创建画布
        self.canvas = tk.Canvas(
            canvas_border,
            width=self.canvas_size,
            height=self.canvas_size,
            bg=self.colors["light"],
            highlightthickness=0,
        )
        self.canvas.pack()

        # 创建PIL图像
        self.image = Image.new("L", (self.canvas_size, self.canvas_size), "white")
        self.draw = ImageDraw.Draw(self.image)

        # 绑定鼠标事件
        self.canvas.bind("<Button-1>", self.start_drawing)
        self.canvas.bind("<B1-Motion>", self.draw_line)
        self.canvas.bind("<ButtonRelease-1>", self.stop_drawing)

        # 绑定快捷键
        self.root.bind("<Control-z>", lambda e: self.undo())
        self.root.bind("<Return>", lambda e: self.predict())
        self.root.bind("<Control-l>", lambda e: self.clear_last_stroke())
        self.root.bind("<Escape>", lambda e: self.clear())

    def update_buttons_state(self, state="disabled"):
        """更新所有操作按钮和画布的状态"""
        def update_frame_buttons(frame):
            """递归更新框架中的按钮状态"""
            for child in frame.winfo_children():
                if isinstance(child, (ttk.Frame, ttk.LabelFrame)):
                    # 递归处理子框架
                    update_frame_buttons(child)
                elif isinstance(child, ttk.Button):
                    # 模型加载相关的按钮始终保持启用
                    if child.cget("text") not in ["选择模型", "加载模型"]:
                        child.configure(state=state)
        
        # 从主框架开始递归更新
        for frame in self.root.winfo_children():
            if isinstance(frame, (ttk.Frame, ttk.LabelFrame)):
                update_frame_buttons(frame)
        
        # 更新画布状态和光标
        if state == "disabled":
            self.canvas.config(cursor="no")  # 禁止光标
            self.canvas.unbind("<Button-1>")
            self.canvas.unbind("<B1-Motion>")
            self.canvas.unbind("<ButtonRelease-1>")
        else:
            self.canvas.config(cursor="crosshair")  # 十字光标
            self.canvas.bind("<Button-1>", self.start_drawing)
            self.canvas.bind("<B1-Motion>", self.draw_line)
            self.canvas.bind("<ButtonRelease-1>", self.stop_drawing)


if __name__ == "__main__":
    app = DigitRecognizer()
    app.run()
